package com.ll.agent.ASMStudy.generate;

import org.objectweb.asm.*;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;

/**
 * @description:
 * @author: LiuLei
 * @create: 2024-07-25
 **/


public class ASMGenerator {
    private static final String CLASS_NAME = "MyClassGenerated";
    private static final String SUPER_CLASS_NAME = "MyClass";

    public static void main(String[] args) throws Exception {
        byte[] classData = generateClass();
        Class<?> generatedClass = new MyClassLoader().defineClass(CLASS_NAME, classData);

        Object instance = generatedClass.getDeclaredConstructor().newInstance();
        MethodHandlers handlers = createMethodHandlers(generatedClass);

        for (int i = 0; i < 100000; i++) {
            long start = System.currentTimeMillis();
            handlers.methods[0].invoke(instance);
            handlers.methods[1].invoke(instance);
            long end = System.currentTimeMillis();
            System.out.println("Time taken: " + (end - start) + "ms");
        }
    }

    // Generate the class bytecode dynamically using ASM
    private static byte[] generateClass() {
        ClassWriter cw = new ClassWriter(0);
        String internalName = CLASS_NAME.replace('.', '/');
        String superInternalName = SUPER_CLASS_NAME.replace('.', '/');

        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, internalName, null, superInternalName, null);

        // Add constructor
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, superInternalName, "<init>", "()V", false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(1, 1);
        mv.visitEnd();
        cw.visitEnd();

        return cw.toByteArray();
    }

    // Create MethodHandlers to optimize method invocation
    private static MethodHandlers createMethodHandlers(Class<?> clazz) {
        Method[] methods = clazz.getMethods();
        MethodHandlers handlers = new MethodHandlers(clazz);

        for (int i = 0; i < methods.length; i++) {
            handlers.methodIndexMap.put(methods[i].getName(), i);
            handlers.methods[i] = methods[i];
        }

        return handlers;
    }

    public static class MethodHandlers {
        public final Map<String, Integer> methodIndexMap = new HashMap<>();
        public final Method[] methods;

        public MethodHandlers(Class<?> clazz) {
            methods = clazz.getMethods();
        }
    }

    public static class MyClassLoader extends ClassLoader {
        public Class<?> defineClass(String name, byte[] b) {
            return defineClass(name, b, 0, b.length);
        }
    }
}

